#!/usr/bin/env python3
"""
MPS GPU 快速启动脚本 - 超小规模训练测试
使用极少量数据和少量epoch，快速验证整个训练流程
"""
import os
import sys
import argparse
import torch
from pathlib import Path

# 添加项目路径
sys.path.append(str(Path(__file__).parent.parent))

def get_device():
    """获取最佳设备"""
    if torch.backends.mps.is_available():
        device = torch.device("mps")
        print(f"✅ 使用 MPS 设备: {device}")
    else:
        device = torch.device("cpu")  
        print(f"⚠️  MPS 不可用，使用 CPU: {device}")
    return device

def quick_cnclip_test(device, data_scale=0.001, max_epochs=1):
    """快速 CN-CLIP 训练测试"""
    print(f"\n=== CN-CLIP 快速训练测试 ===")
    print(f"数据规模: {data_scale*100}%")
    print(f"训练轮数: {max_epochs}")
    print(f"设备: {device}")
    
    try:
        import pytorch_lightning as pl
        from pytorch_lightning.callbacks import RichProgressBar
        from pytorch_lightning.loggers import TensorBoardLogger
        
        # 检查数据集配置
        config_path = Path(__file__).parent.parent / "dataset" / "config.json"
        if not config_path.exists():
            print("❌ 数据集配置文件不存在，请先配置数据集路径")
            print(f"需要创建: {config_path}")
            return False
            
        # 导入模型和数据
        from cnclip_finetune.model import CNClipLightning, create_data_module
        
        # 创建数据模块 - 超小规模
        data_module = create_data_module(
            data_scale=data_scale,
            batch_size=2,  # 极小批次
            num_workers=0  # MPS 建议不使用多进程
        )
        
        # 创建模型
        model = CNClipLightning(
            model_name="ViT-B-16",
            learning_rate=1e-5,
            weight_decay=0.01,
            device_type=str(device)
        )
        
        # 配置训练器 - 快速测试设置
        trainer = pl.Trainer(
            max_epochs=max_epochs,
            accelerator="auto",
            devices=1,
            logger=TensorBoardLogger("codebase/mps_test/logs", name="quick_test"),
            callbacks=[RichProgressBar()],
            fast_dev_run=False,  # 运行完整但极短的训练
            limit_train_batches=5,  # 只训练5个批次
            limit_val_batches=2,    # 只验证2个批次  
            enable_checkpointing=False,  # 不保存检查点
            log_every_n_steps=1
        )
        
        print("开始快速训练测试...")
        trainer.fit(model, data_module)
        
        print("✅ 快速训练测试完成！")
        return True
        
    except Exception as e:
        print(f"❌ 快速训练失败: {e}")
        import traceback
        traceback.print_exc()
        return False

def simple_inference_test(device):
    """简单推理测试"""
    print(f"\n=== 简单推理测试 (设备: {device}) ===")
    
    try:
        import cn_clip.clip as clip
        from cn_clip.clip import load_from_name, tokenize
        
        # 加载预训练模型
        model, preprocess = load_from_name("ViT-B-16", device=device, download_root="./assets/checkpoints")
        model.eval()
        
        # 创建测试数据
        image = torch.randn(1, 3, 224, 224).to(device)
        texts = ["文物", "古代艺术品", "青铜器", "瓷器"]
        text_tokens = tokenize(texts).to(device)
        
        # 推理
        with torch.no_grad():
            logits_per_image, logits_per_text = model(image, text_tokens)
            probs = logits_per_image.softmax(dim=-1).cpu().numpy()
            
        print("文本匹配概率:")
        for text, prob in zip(texts, probs[0]):
            print(f"  {text}: {prob:.3f}")
            
        print("✅ 推理测试成功")
        return True
        
    except Exception as e:
        print(f"❌ 推理测试失败: {e}")
        return False

def main():
    """主函数"""
    parser = argparse.ArgumentParser(description="MPS GPU 快速测试")
    parser.add_argument("--data-scale", type=float, default=0.001, 
                       help="数据规模比例 (默认: 0.001 = 0.1%)")
    parser.add_argument("--max-epochs", type=int, default=1,
                       help="最大训练轮数 (默认: 1)")
    parser.add_argument("--inference-only", action="store_true",
                       help="仅运行推理测试")
    
    args = parser.parse_args()
    
    print("🚀 WenwuClip MPS GPU 快速测试")
    print("=" * 50)
    
    # 获取设备
    device = get_device()
    
    success = True
    
    # 推理测试
    if not simple_inference_test(device):
        success = False
    
    # 训练测试 (如果不是仅推理模式)
    if not args.inference_only:
        if not quick_cnclip_test(device, args.data_scale, args.max_epochs):
            success = False
    
    print("\n" + "=" * 50)
    if success:
        print("🎉 所有测试通过！MPS 设备可以正常使用")
        print("\n📝 接下来可以尝试:")
        print("1. 配置你的数据集路径 (codebase/dataset/config.json)")
        print("2. 运行更大规模的训练:")
        print("   python codebase/cnclip_finetune/train.py --data-scale=0.1 --max-epochs=3")
    else:
        print("⚠️  部分测试失败，请检查错误信息")

if __name__ == "__main__":
    main()